import matplotlib.pyplot as plt
import os

# 读取文本文件并解析日志数据
log_file = "Low_Contrast/UNet-40epoch-ReduceLROnPlateau/valid_log.txt"  # 请替换为你的日志文件路径
with open(log_file, 'r') as f:
    lines = f.readlines()

epochs = []
losses = []
accuracies = []
acc_classes = []
mious = []
ious = []
precisions = []
recalls = []
f1s = []

for line in lines:
    if "Epoch" in line:
        parts = line.split(', ')
        epoch = int(parts[0].split('[')[1].split('/')[0])
        loss = float(parts[2].split(' ')[1])
        accuracy = float(parts[3].split(' ')[1])
        acc_class = float(parts[4].split(' ')[1])
        mIOU = float(parts[5].split(' ')[1])
        IOU = float(parts[6].split(' ')[1])
        precision = float(parts[7].split(' ')[1])
        recall = float(parts[8].split(' ')[1])
        f1 = float(parts[9].split(' ')[1])

        epochs.append(epoch)
        losses.append(loss)
        accuracies.append(accuracy)
        acc_classes.append(acc_class)
        mious.append(mIOU)
        ious.append(IOU)
        precisions.append(precision)
        recalls.append(recall)
        f1s.append(f1)


def plot_and_save(epochs, losses, accuracies, mious, ious, precisions, recalls, f1s, folder_path):
    # 创建images文件夹
    os.makedirs("images", exist_ok=True)

    # 绘制损失图
    plt.figure(figsize=(6, 4))
    plt.plot(epochs, losses, marker='o', color='b')
    plt.title('Loss over Epochs')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.savefig(f"images/{folder_path}-Loss.png")
    plt.close()

    # 绘制准确率和IOU图
    plt.figure(figsize=(6, 4))
    plt.plot(epochs, accuracies, marker='o', color='r')
    plt.plot(epochs, mious, marker='o', color='g')
    plt.plot(epochs, ious, marker='o', color='m')
    plt.legend(['Accuracy', 'mIOU', 'IOU'])
    plt.title('Accuracy and IOU over Epochs')
    plt.xlabel('Epoch')
    plt.ylabel('Value')
    plt.savefig(f"images/{folder_path}-Accuracy-IOU.png")
    plt.close()

    # 绘制精度、召回率和F1图
    plt.figure(figsize=(6, 4))
    plt.plot(epochs, precisions, marker='o', color='c')
    plt.plot(epochs, recalls, marker='o', color='y')
    plt.plot(epochs, f1s, marker='o', color='k')
    plt.legend(['Precision', 'Recall', 'F1'])
    plt.title('Precision, Recall, and F1 over Epochs')
    plt.xlabel('Epoch')
    plt.ylabel('Value')
    plt.savefig(f"images/{folder_path}-Precision-Recall-F1.png")
    plt.close()


# 使用函数
folder_path = log_file.split('/')[1]  # 例如："Low_Contrast/SegNet"
print(folder_path)
plot_and_save(epochs, losses, accuracies, mious, ious, precisions, recalls, f1s, folder_path)
